#include "barrier.hpp"
#include "threadpool.h"
#include <iostream>
#define BLOCK_SIZE 32
#define __syncthreads() barrier.wait();
void AddKernel(int blockIdx, int threadIdx, int gridDim, int blockDim,
               float shared_a[BLOCK_SIZE], float shared_b[BLOCK_SIZE],
               Barrier &barrier, float *a, float *b, float *c, int N) {
    int tid = threadIdx + blockIdx * blockDim;
    if (tid < N) {

        shared_a[threadIdx] = a[tid];
        shared_b[threadIdx] = b[tid];
    }
    __syncthreads();
    if (tid < N) {
        c[tid] = shared_a[threadIdx] + shared_b[threadIdx];
    }
}

void LaunchKernel(int gridDim, int blockDim, float *a, float *b, float *c,
                  int N) {
    ThreadPool grid_mananger(gridDim);
    std::vector<std::future<void>> grid_results;
    for (int blockIdx = 0; blockIdx < gridDim; ++blockIdx) {
        grid_results.emplace_back(grid_mananger.enqueue([=] {
            ThreadPool block_manager(blockDim);
            std::vector<std::future<void>> block_results;
            float shared_a[BLOCK_SIZE];
            float shared_b[BLOCK_SIZE];
            Barrier barrier(blockDim);
            for (int threadIdx = 0; threadIdx < blockDim; ++threadIdx) {
                block_manager.enqueue([=, &shared_a, &shared_b, &barrier] {
                    AddKernel(blockIdx, threadIdx, gridDim, blockDim, shared_a,
                              shared_b, barrier, a, b, c, N);
                });
            }
            for (auto &&result : block_results)
                result.get();
        }));
    }

    for (auto &&result : grid_results)
        result.get();
}

int main() {
    int N = 100;
    int blockDim = BLOCK_SIZE;
    int gridDim = (N + blockDim - 1) / blockDim;
    auto a = new float[N];
    auto b = new float[N];
    auto c = new float[N];
    for (int i = 0; i < N; i++) {
        a[i] = i;
        b[i] = i;
        c[i] = 0;
    }
    LaunchKernel(gridDim, blockDim, a, b, c, N);
    for (int i = 0; i < N; i++) {
        std::cout << c[i] << " ";
    }
    std::cout << std::endl;

    delete[] a;
    delete[] b;
    delete[] c;
    return 0;
}